{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/dtypes.py:516: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/dtypes.py:517: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/dtypes.py:518: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/dtypes.py:519: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/dtypes.py:520: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorflow/python/framework/dtypes.py:525: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorboard/compat/tensorflow_stub/dtypes.py:541: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint8 = np.dtype([(\"qint8\", np.int8, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorboard/compat/tensorflow_stub/dtypes.py:542: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint8 = np.dtype([(\"quint8\", np.uint8, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorboard/compat/tensorflow_stub/dtypes.py:543: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint16 = np.dtype([(\"qint16\", np.int16, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorboard/compat/tensorflow_stub/dtypes.py:544: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_quint16 = np.dtype([(\"quint16\", np.uint16, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorboard/compat/tensorflow_stub/dtypes.py:545: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  _np_qint32 = np.dtype([(\"qint32\", np.int32, 1)])\n",
      "/usr/local/lib/python3.6/dist-packages/tensorboard/compat/tensorflow_stub/dtypes.py:550: FutureWarning: Passing (type, 1) or '1type' as a synonym of type is deprecated; in a future version of numpy, it will be understood as (type, (1,)) / '(1,)type'.\n",
      "  np_resource = np.dtype([(\"resource\", np.ubyte, 1)])\n"
     ]
    }
   ],
   "source": [
    "\"\"\"\n",
    "solving pendulum using actor-critic model\n",
    "\"\"\"\n",
    "\n",
    "import gym\n",
    "import numpy as np\n",
    "from tensorflow.keras.models import Sequential, Model\n",
    "from tensorflow.keras.layers import Dense, Dropout, Input, GRU, Masking\n",
    "from tensorflow.keras.layers import Add, Concatenate\n",
    "from tensorflow.keras.optimizers import Adam\n",
    "import tensorflow.keras.backend as K\n",
    "\n",
    "import tensorflow as tf\n",
    "\n",
    "import random\n",
    "from collections import deque\n",
    "\n",
    "def stack_samples(samples):\n",
    "\tarray = np.array(samples)\n",
    "\t\n",
    "\tcurrent_states = np.stack(array[:,0]).reshape((array.shape[0],-1,array[0,0].shape[2]))\n",
    "\tactions = np.stack(array[:,1]).reshape((array.shape[0],-1))\n",
    "\trewards = np.stack(array[:,2]).reshape((array.shape[0],-1))\n",
    "\tnew_states = np.stack(array[:,3]).reshape((array.shape[0],-1,array[0,3].shape[2]))\n",
    "\tdones = np.stack(array[:,4]).reshape((array.shape[0],1))\n",
    "\t\n",
    "\treturn current_states, actions, rewards, new_states, dones\n",
    "\n",
    "# determines how to assign values to each state, i.e. takes the state\n",
    "# and action (two-input model) and determines the corresponding value\n",
    "class ActorCritic:\n",
    "\tdef __init__(self, env, sess):\n",
    "\t\tself.env  = env\n",
    "\t\tself.sess = sess\n",
    "\n",
    "\t\tself.learning_rate = 0.0001\n",
    "\t\tself.epsilon = .9\n",
    "\t\tself.epsilon_decay = .99995\n",
    "\t\tself.gamma = .90\n",
    "\t\tself.tau   = .01\n",
    "\n",
    "\t\t# ===================================================================== #\n",
    "\t\t#                               Actor Model                             #\n",
    "\t\t# Chain rule: find the gradient of chaging the actor network params in  #\n",
    "\t\t# getting closest to the final value network predictions, i.e. de/dA    #\n",
    "\t\t# Calculate de/dA as = de/dC * dC/dA, where e is error, C critic, A act #\n",
    "\t\t# ===================================================================== #\n",
    "\t\tself.ema = tf.train.ExponentialMovingAverage(decay=1-self.tau)\n",
    "\t\tself.memory = deque(maxlen=4000)\n",
    "\t\tself.actor_state_input, self.actor_model = self.create_actor_model()\n",
    "\t\t_, self.target_actor_model = self.create_actor_model()\n",
    "\n",
    "\t\tself.actor_critic_grad = tf.placeholder(tf.float32,\n",
    "\t\t\t[None, self.env.action_space.shape[0]]) # where we will feed de/dC (from critic)\n",
    "\n",
    "\t\tactor_model_weights = self.actor_model.trainable_weights\n",
    "\t\tself.actor_grads = tf.gradients(self.actor_model.output,\n",
    "\t\t\tactor_model_weights, -self.actor_critic_grad) # dC/dA (from actor)\n",
    "\t\tgrads = zip(self.actor_grads, actor_model_weights)\n",
    "\t\tself.optimize = tf.train.AdamOptimizer(self.learning_rate).apply_gradients(grads)\n",
    "\n",
    "\t\t# ===================================================================== #\n",
    "\t\t#                              Critic Model                             #\n",
    "\t\t# ===================================================================== #\n",
    "\n",
    "\t\tself.critic_state_input, self.critic_action_input, \\\n",
    "\t\t\tself.critic_model = self.create_critic_model()\n",
    "\t\t_, _, self.target_critic_model = self.create_critic_model()\n",
    "\n",
    "\t\tself.critic_grads = tf.gradients(self.critic_model.output,\n",
    "\t\t\tself.critic_action_input) # where we calcaulte de/dC for feeding above\n",
    "\n",
    "\t\t# Initialize for later gradient calculations\n",
    "\t\tself.sess.run(tf.initialize_all_variables())\n",
    "\n",
    "\t# ========================================================================= #\n",
    "\t#                              Model Definitions                            #\n",
    "\n",
    "\tdef create_actor_model(self):\n",
    "\t# ========================================================================= #\n",
    "\t\tstate_input = Input(shape=(None,self.env.observation_space.shape[0]))\n",
    "\t\tmask_state_input = Masking(mask_value=0.)(state_input)\n",
    "\t\th1 = Dense(500, activation='relu')(mask_state_input)\n",
    "\t\tactor_rnn,state_h = GRU(256, return_state=True)(h1)\n",
    "\t\th2 = Dense(500, activation='relu')(state_h)\n",
    "\t\toutput = Dense(self.env.action_space.shape[0], activation='tanh')(h2)\n",
    "\n",
    "\t\tmodel = Model([state_input], output)\n",
    "\t\tadam  = Adam(lr=0.0001)\n",
    "\t\tmodel.compile(loss=\"mse\", optimizer=adam)\n",
    "\t\treturn state_input, model\n",
    "\n",
    "\tdef create_critic_model(self):\n",
    "\t\tstate_input = Input(shape=(None,self.env.observation_space.shape[0]))\n",
    "\t\tmask_state_input = Masking(mask_value=0.)(state_input)\n",
    "\t\tstate_h1 = Dense(500, activation='relu')(mask_state_input)\n",
    "\t\tcritic_rnn,state_h2 = GRU(256, return_state=True)(state_h1)\n",
    "\n",
    "\t\taction_input = Input(shape=self.env.action_space.shape)\n",
    "\t\taction_h1    = Dense(500)(action_input)\n",
    "\n",
    "\t\tmerged    = Concatenate()([state_h2, action_h1])\n",
    "\t\tmerged_h1 = Dense(500, activation='relu')(merged)\n",
    "\t\toutput = Dense(1, activation='linear')(merged_h1)\n",
    "\t\tmodel  = Model([state_input,action_input],output)\n",
    "\n",
    "\t\tadam  = Adam(lr=0.0001)\n",
    "\t\tmodel.compile(loss=\"mse\", optimizer=adam)\n",
    "\t\treturn state_input, action_input, model\n",
    "\n",
    "\t# ========================================================================= #\n",
    "\t#                               Model Training                              #\n",
    "\t# ========================================================================= #\n",
    "\n",
    "\tdef remember(self, cur_state, action, reward, new_state, done):\n",
    "\t\tself.memory.append([cur_state, action, reward, new_state, done])\n",
    "\n",
    "\tdef _train_actor(self, samples):\n",
    "\t\t\n",
    "\t\t\tcur_states, actions, rewards, new_states, _ =  stack_samples(samples)\n",
    "\t\t\tpredicted_actions = self.actor_model.predict(cur_states)\n",
    "\t\t\tgrads = self.sess.run(self.critic_grads, feed_dict={\n",
    "\t\t\t\tself.critic_state_input:  cur_states,\n",
    "\t\t\t\tself.critic_action_input: predicted_actions\n",
    "\t\t\t})[0]\n",
    "\n",
    "\t\t\tself.sess.run(self.optimize, feed_dict={\n",
    "\t\t\t\tself.actor_state_input: cur_states,\n",
    "\t\t\t\tself.actor_critic_grad: grads\n",
    "\t\t\t})\n",
    "\n",
    "\tdef _train_critic(self, samples):\n",
    "   \n",
    "\n",
    "\t\tcur_states, actions, rewards, new_states, dones = stack_samples(samples)\n",
    "\t\ttarget_actions = self.target_actor_model.predict(new_states)\n",
    "\t\tfuture_rewards = self.target_critic_model.predict([new_states, target_actions])\n",
    "\t\tdones = dones.reshape(rewards.shape)\n",
    "\t\tfuture_rewards = future_rewards.reshape(rewards.shape)\n",
    "\t\trewards += self.gamma * future_rewards * (1 - dones)\n",
    "\t\t\n",
    "\t\tevaluation = self.critic_model.fit([cur_states, actions], rewards, verbose=0)\n",
    "\t\t#print(evaluation.history)\n",
    "        \n",
    "\tdef train(self):\n",
    "\t\tbatch_size = 256\n",
    "\t\tif len(self.memory) < batch_size:\n",
    "\t\t\treturn\n",
    "\n",
    "\t\trewards = []\n",
    "\t\tsamples = random.sample(self.memory, batch_size)\n",
    "\t\tself.samples = samples\n",
    "\t\tself._train_critic(samples)\n",
    "\t\tself._train_actor(samples)\n",
    "\n",
    "\t# ========================================================================= #\n",
    "\t#                         Target Model Updating                             #\n",
    "\t# ========================================================================= #\n",
    "\n",
    "\tdef _update_actor_target(self):\n",
    "\t\tactor_model_weights  = self.actor_model.get_weights()\n",
    "\t\tactor_target_weights = self.target_actor_model.get_weights()\n",
    "\t\tsess.run(ema.apply(actor_model_weights))\n",
    "\t\tself.target_actor_model.set_weights(ema.average(actor_model_weights))\n",
    "\n",
    "\tdef _update_critic_target(self):\n",
    "\t\tcritic_model_weights  = self.critic_model.get_weights()\n",
    "\t\tcritic_target_weights = self.target_critic_model.get_weights()\n",
    "\t\tsess.run(ema.apply(critic_model_weights))\n",
    "\t\tself.target_actor_model.set_weights(ema.average(critic_model_weights))\n",
    "\n",
    "\tdef update_target(self):\n",
    "\t\tself._update_actor_target()\n",
    "\t\tself._update_critic_target()\n",
    "\n",
    "\t# ========================================================================= #\n",
    "\t#                              Model Predictions                            #\n",
    "\t# ========================================================================= #\n",
    "\n",
    "\tdef act(self, cur_state):\n",
    "\t\tself.epsilon *= self.epsilon_decay\n",
    "\t\tif np.random.random() < self.epsilon:\n",
    "\t\t\treturn self.actor_model.predict(cur_state)*2 + np.random.normal()\n",
    "\t\treturn self.actor_model.predict(cur_state)*2\n",
    "\n",
    "\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/ops/init_ops.py:1251: calling VarianceScaling.__init__ (from tensorflow.python.ops.init_ops) with dtype is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Call initializer instance with the dtype argument instead of passing it to the constructor\n"
     ]
    },
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/usr/local/lib/python3.6/dist-packages/gym/logger.py:30: UserWarning: \u001b[33mWARN: Box bound precision lowered by casting to float32\u001b[0m\n",
      "  warnings.warn(colorize('%s: %s'%('WARN', msg % args), 'yellow'))\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/backend.py:3794: add_dispatch_support.<locals>.wrapper (from tensorflow.python.ops.array_ops) is deprecated and will be removed in a future version.\n",
      "Instructions for updating:\n",
      "Use tf.where in 2.0, which has the same broadcast rule as np.where\n",
      "WARNING:tensorflow:From /usr/local/lib/python3.6/dist-packages/tensorflow/python/util/tf_should_use.py:193: initialize_all_variables (from tensorflow.python.ops.variables) is deprecated and will be removed after 2017-03-02.\n",
      "Instructions for updating:\n",
      "Use `tf.global_variables_initializer` instead.\n",
      "trial:0\n",
      "[-18.885056]\n",
      "trial:1\n",
      "[-14.946356]\n",
      "trial:2\n",
      "[-23.934801]\n",
      "trial:3\n",
      "[-13.955161]\n",
      "trial:4\n",
      "[-13.828778]\n",
      "trial:5\n",
      "[-16.940048]\n",
      "trial:6\n",
      "[-14.42596]\n",
      "trial:7\n",
      "[-13.675378]\n",
      "trial:8\n",
      "[-25.353037]\n",
      "trial:9\n",
      "[-5.213123]\n",
      "trial:10\n",
      "[-6.591884]\n",
      "trial:11\n",
      "[-0.6047917]\n",
      "trial:12\n",
      "[-11.55867]\n",
      "trial:13\n",
      "[-23.064182]\n",
      "trial:14\n",
      "[-7.3252697]\n",
      "trial:15\n",
      "[-14.480252]\n",
      "trial:16\n",
      "[-29.000439]\n",
      "trial:17\n",
      "[-0.96139324]\n",
      "trial:18\n",
      "[-22.788048]\n",
      "trial:19\n",
      "[-25.671108]\n",
      "trial:20\n",
      "[-32.3016]\n",
      "trial:21\n",
      "[-7.8974853]\n",
      "trial:22\n",
      "[-25.5116]\n",
      "trial:23\n",
      "[-10.96522]\n",
      "trial:24\n",
      "[-11.30288]\n",
      "trial:25\n",
      "[-16.254215]\n",
      "trial:26\n",
      "[-10.478016]\n",
      "trial:27\n",
      "[-19.02015]\n",
      "trial:28\n",
      "[-7.5316234]\n",
      "trial:29\n",
      "[-24.208727]\n",
      "trial:30\n",
      "[-22.037844]\n",
      "trial:31\n",
      "[-15.894063]\n",
      "trial:32\n",
      "[-19.221386]\n",
      "trial:33\n",
      "[-8.053923]\n",
      "trial:34\n",
      "[-23.036993]\n",
      "trial:35\n",
      "[-20.32407]\n",
      "trial:36\n",
      "[-19.289873]\n",
      "trial:37\n",
      "[-23.3839]\n",
      "trial:38\n",
      "[-24.926517]\n",
      "trial:39\n",
      "[-21.299429]\n",
      "trial:40\n",
      "[-17.683779]\n",
      "trial:41\n",
      "[-23.840548]\n",
      "trial:42\n",
      "[-9.397272]\n",
      "trial:43\n",
      "[-14.794532]\n",
      "trial:44\n",
      "[-18.294945]\n",
      "trial:45\n",
      "[-12.407223]\n",
      "trial:46\n",
      "[-27.698185]\n",
      "trial:47\n",
      "[-23.095171]\n",
      "trial:48\n"
     ]
    }
   ],
   "source": [
    "def main():\n",
    "\tsess = tf.Session()\n",
    "\tK.set_session(sess)\n",
    "\tenv = gym.make(\"Pendulum-v0\")\n",
    "\tactor_critic = ActorCritic(env, sess)\n",
    "\n",
    "\tnum_trials = 10000\n",
    "\ttrial_len  = 200\n",
    "\tmax_seq = 1\n",
    "\n",
    "\tfor i in range(num_trials):\n",
    "\t\tprint(\"trial:\" + str(i))\n",
    "\t\tcur_state = env.reset()\n",
    "\t\taction = env.action_space.sample()\n",
    "\t\treward_sum = 0\n",
    "\t\tobs_list = []        \n",
    "\t\tobs_list.append(cur_state)\n",
    "\t\tfor j in range(trial_len):\n",
    "\t\t\t#env.render()        \n",
    "\t\t\tobs_seq = np.asarray(obs_list)\n",
    "\t\t\tobs_seq = obs_seq.reshape((1, -1, env.observation_space.shape[0]))\n",
    "\t\t\taction = actor_critic.act(obs_seq)\n",
    "\t\t\taction = action.reshape((1, env.action_space.shape[0]))\n",
    "\n",
    "\t\t\tnew_state, reward, done, _ = env.step(action)\n",
    "\t\t\treward += reward\n",
    "\t\t\tif j == (trial_len - 1):\n",
    "\t\t\t\tdone = True\n",
    "\t\t\t\tprint(reward)\n",
    "\n",
    "\t\t\tif (j % 5 == 0):\n",
    "\t\t\t\tactor_critic.train()\n",
    "\t\t\t\tactor_critic.update_target()   \n",
    "\t\t\t\n",
    "\t\t\tnew_state = new_state.reshape((env.observation_space.shape))\n",
    "\n",
    "\t\t\tobs_list.append(new_state)\n",
    "\t\t\tif len(obs_list)>max_seq:\n",
    "\t\t\t    obs_list.pop(0)\n",
    "\t\t\tnext_obs_seq = np.asarray(obs_list)\n",
    "\t\t\tnext_obs_seq = next_obs_seq.reshape((1, -1, env.observation_space.shape[0]))\n",
    "            \n",
    "\t\t\t#padding\n",
    "\t\t\tpad_width = max_seq-np.size(obs_seq,1)\n",
    "\t\t\tobs_seq = np.pad(obs_seq,((0,0),(pad_width,0),(0,0)))\n",
    "\t\t\tnext_obs_seq = np.pad(next_obs_seq,((0,0),(pad_width,0),(0,0)))\n",
    "\t\t\t#print(obs_seq.shape)\n",
    "\t\t\t#print(next_obs_seq.shape)\n",
    "            if obs_list\n",
    "            \n",
    "\t\t\tactor_critic.remember(obs_seq, action, reward, next_obs_seq, done)\n",
    "\t\t\tcur_state = new_state\n",
    "\n",
    "\t\tif (i % 5 == 0) and i!=0:\n",
    "\t\t\tcur_state = env.reset()\n",
    "\t\t\tobs_list = []                   \n",
    "\t\t\tfor j in range(500):\n",
    "\t\t\t\t#env.render()\n",
    "\t\t\t\tcur_state = cur_state.reshape((1, env.observation_space.shape[0]))\n",
    "\t\t\t\tobs_list.append(cur_state)\n",
    "\t\t\t\tobs_seq = np.asarray(obs_list)\n",
    "\t\t\t\tobs_seq = obs_seq.reshape((1, -1, env.observation_space.shape[0]))\n",
    "\t\t\t\taction = actor_critic.act(obs_seq)\n",
    "\t\t\t\taction = action.reshape((1, env.action_space.shape[0]))\n",
    "\n",
    "\t\t\t\tnew_state, reward, done, _ = env.step(action)\n",
    "\t\t\t\t#reward += reward\n",
    "\t\t\t\t#if j == (trial_len - 1):\n",
    "\t\t\t\t\t#done = True\n",
    "\t\t\t\t\t#print(reward)\n",
    "\n",
    "\t\t\t\t#if (j % 5 == 0):\n",
    "\t\t\t\t#    actor_critic.train()\n",
    "\t\t\t\t#    actor_critic.update_target()   \n",
    "\t\t\t\t\n",
    "\t\t\t\tnew_state = new_state.reshape((1, env.observation_space.shape[0]))\n",
    "\n",
    "\t\t\t\t#actor_critic.remember(cur_state, action, reward, new_state, done)\n",
    "\t\t\t\tcur_state = new_state\n",
    "\t\t\t\t\n",
    "\n",
    "if __name__ == \"__main__\":\n",
    "\tmain()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.6.9"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
